from ultralytics import YOLO
import numpy as np
from typing import List, Tuple, Optional
import cv2

class AIDetector:
    def __init__(self, model_path: str, conf_threshold: float = 0.5):
        self.model = YOLO(model_path)
        self.conf_threshold = conf_threshold
        self.is_running = False

    def detect(self, frame: np.ndarray) -> Tuple[np.ndarray, List[dict]]:
        """
        对图像进行目标检测
        返回: (处理后的图像, 检测结果列表)
        """
        if not self.is_running:
            return frame, []

        results = self.model(frame, conf=self.conf_threshold)
        annotated_frame = results[0].plot()
        
        detections = []
        for r in results[0].boxes.data:
            x1, y1, x2, y2, conf, cls = r
            detections.append({
                'bbox': [float(x1), float(y1), float(x2), float(y2)],
                'confidence': float(conf),
                'class': int(cls)
            })

        return annotated_frame, detections

    def start(self):
        self.is_running = True

    def stop(self):
        self.is_running = False 